{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "68e1c158",
   "metadata": {},
   "source": [
    "# 通过SK实现RAG\n",
    "\n",
    "对于需要外部知识支撑的场景，我们通常会使用RAG（Retrieval Augmented Generation）来实现，其中一般会涉及到文档解析，切片，向量化\n",
    "检索，通过LLM生成输出等步骤。\n",
    "在[Semantic ChatBot](./chatbot_with_sk.ipynb)中我们基于semantic kernel实现了一个简单的chatbot，并使用`context variables`\n",
    "实现了历史聊天的存储和记录。\n",
    "但是对于很长的外部知识库库，我们可能会因此需要超长的上下文，甚至无法没办法记录所有的历史，为此SK中提供了Memory类型以记录实现LLM的记忆能力。"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "890c58b870d009b0",
   "metadata": {},
   "outputs": [],
   "source": [
    "! pip install semantic-kernel==0.4.5.dev0\n",
    "! pip install chromadb==0.4.14"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "851c62ca",
   "metadata": {},
   "source": [
    "初始化鉴权，导入SK，以及适配SK的Qianfan实现类型"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "id": "5f6cd63f",
   "metadata": {},
   "outputs": [],
   "source": [
    "import os \n",
    "\n",
    "os.environ[\"QIANFAN_ACCESS_KEY\"] = \"your_ak\"\n",
    "os.environ[\"QIANFAN_SECRET_KEY\"] = \"your_sk\""
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "id": "508ad44f",
   "metadata": {},
   "outputs": [],
   "source": [
    "import semantic_kernel as sk\n",
    "from qianfan.extensions.semantic_kernel import (\n",
    "    QianfanChatCompletion,\n",
    "    QianfanTextEmbedding,\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "d8ddffc1",
   "metadata": {},
   "source": [
    "### SK Memory\n",
    "\n",
    "SK Memory 是一个数据框架，可以通过接入外部的各种数据源；可以是从网页，数据库，email等，这些都集成在了SK的内置connectors中，而通过`QianfanTextEmbedding`，可以提取这些数据源中的文本的特征向量，以供后续的检索使用。\n",
    "\n",
    "\n",
    "这里使用了`VolatileMemoryStore`作为Memory的实现为例，`VolatileMemoryStore` 实现了内存的临时存储（底层通过一个Dict[Dict[str, MemoryRecord]] 实现分collection的kv存储）。"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "id": "8f8dcbc6",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "{'recall': SKFunction(), 'save': SKFunction()}"
      ]
     },
     "execution_count": 3,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "from semantic_kernel.memory import VolatileMemoryStore\n",
    "from semantic_kernel.core_skills import TextMemorySkill\n",
    "kernel = sk.Kernel()\n",
    "\n",
    "qf_chat_service = QianfanChatCompletion(ai_model_id=\"ERNIE-Bot-4\")\n",
    "qf_text_embedding = QianfanTextEmbedding(ai_model_id=\"Embedding-V1\")\n",
    "kernel.add_chat_service(\"chat-qf\", qf_chat_service)\n",
    "kernel.add_text_embedding_generation_service(\"embed-eb\", qf_text_embedding)\n",
    "\n",
    "kernel.register_memory_store(memory_store=VolatileMemoryStore())\n",
    "kernel.import_skill(TextMemorySkill())"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "65fbedd3",
   "metadata": {},
   "source": [
    "调用异步函数，完成数据的添加，这里往了一个名为`aboutMe`的`collection`中添加了若干个人信息"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "id": "d096504c",
   "metadata": {},
   "outputs": [],
   "source": [
    "async def populate_memory(kernel: sk.Kernel) -> None:\n",
    "    # Add some documents to the semantic memory\n",
    "    await kernel.memory.save_information_async(collection=\"aboutMe\", id=\"info1\", text=\"我名字叫做小度\")\n",
    "    await kernel.memory.save_information_async(\n",
    "        collection=\"aboutMe\", id=\"info2\", text=\"我工作在baidu\"\n",
    "    )\n",
    "    await kernel.memory.save_information_async(\n",
    "        collection=\"aboutMe\", id=\"info3\", text=\"我来自中国\"\n",
    "    )\n",
    "    await kernel.memory.save_information_async(\n",
    "        collection=\"aboutMe\",\n",
    "        id=\"info4\",\n",
    "        text=\"我曾去过北京，上海，深圳\",\n",
    "    )\n",
    "    await kernel.memory.save_information_async(collection=\"aboutMe\", id=\"info5\", text=\"我爱打羽毛球\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "id": "2a87c08e",
   "metadata": {},
   "outputs": [],
   "source": [
    "await populate_memory(kernel)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "2calf857",
   "metadata": {},
   "source": [
    "通过`TextMemoryBase`中实现的余弦相似度计算向量相似度可以search到对应相似的回答："
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "id": "628c843e",
   "metadata": {},
   "outputs": [],
   "source": [
    "async def search_memory_examples(kernel: sk.Kernel) -> None:\n",
    "    questions = [\n",
    "        \"我的名字是？\",\n",
    "        \"我在哪里工作？\",\n",
    "        \"我去过哪些地方旅游?\",\n",
    "        \"我的家乡是?\",\n",
    "        \"我的爱好是？\",\n",
    "    ]\n",
    "\n",
    "    for question in questions:\n",
    "        print(f\"Question: {question}\")\n",
    "        result = await kernel.memory.search_async(\"aboutMe\", question)\n",
    "        print(f\"Answer: {result[0].text}\\n\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "id": "031373f7",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Question: 我的名字是？\n",
      "Answer: 我名字叫做小度\n",
      "\n",
      "Question: 我在哪里工作？\n",
      "Answer: 我工作在baidu\n",
      "\n",
      "Question: 我去过哪些地方旅游?\n",
      "Answer: 我曾去过北京，上海，深圳\n",
      "\n",
      "Question: 我的家乡是?\n",
      "Answer: 我来自中国\n",
      "\n",
      "Question: 我的爱好是？\n",
      "Answer: 我爱打羽毛球\n"
     ]
    }
   ],
   "source": [
    "await search_memory_examples(kernel)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "1ed54a32",
   "metadata": {},
   "source": [
    "### 如果结合对话场景\n",
    "如何将外部知识库和对话系统进行融合？ SK中提供了`TextMemorySkill`，其中包含了`recall`function，可以获取一个input并在kernel的Memory之上执行相似度检索"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "id": "fb8549b2",
   "metadata": {},
   "outputs": [],
   "source": [
    "from typing import Tuple\n",
    "async def setup_chat_with_memory(\n",
    "    kernel: sk.Kernel,\n",
    ") -> Tuple[sk.SKFunctionBase, sk.SKContext]:\n",
    "    from semantic_kernel.core_skills import TextMemorySkill\n",
    "    sk_prompt = \"\"\"\n",
    "    你是一个人设分析师，你需要根据以下背景资料，以及历史聊天记录\n",
    "    回答以下当前输入：\n",
    "    \n",
    "    背景资料：\n",
    "    - {{$fact1}} {{recall $fact1}}\n",
    "    - {{$fact2}} {{recall $fact2}}\n",
    "    - {{$fact3}} {{recall $fact3}}\n",
    "    - {{$fact4}} {{recall $fact4}}\n",
    "    - {{$fact5}} {{recall $fact5}}\n",
    "\n",
    "    聊天记录:\n",
    "    {{$chat_history}}\n",
    "    \n",
    "    当前输入: {{$user_input}}\n",
    "    回答: \n",
    "\n",
    "     \"\"\".strip()\n",
    "\n",
    "    chat_func = kernel.create_semantic_function(sk_prompt, temperature=0.8)\n",
    "\n",
    "    context = kernel.create_new_context()\n",
    "    context[\"fact1\"] = \"我的名字是？\"\n",
    "    context[\"fact2\"] = \"我在哪里工作？\"\n",
    "    context[\"fact3\"] = \"我去过哪些地方旅游?\"\n",
    "    context[\"fact4\"] = \"我的家乡是?\"\n",
    "    context[\"fact5\"] = \"我的爱好是？\"\n",
    "\n",
    "    context[sk.core_skills.TextMemorySkill.COLLECTION_PARAM] = \"aboutMe\"\n",
    "    # context[sk.core_skills.TextMemorySkill.RELEVANCE_PARAM] = \"0.5\"\n",
    "\n",
    "    context[\"chat_history\"] = \"\"\n",
    "    return chat_func, context"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "1ac62457",
   "metadata": {},
   "source": [
    "其中`RelevanceParam` 用于指定检索的阈值，`COLLECTION_PARAM`用于指定collection名，sk_prompt中的recall 是 `TextMemorySkill` 中的一个NativeFunction。相当于`TextMemorySkill.search_async`"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "id": "75267a2f",
   "metadata": {},
   "outputs": [],
   "source": [
    "#-# cell_skip\n",
    "async def chat(kernel: sk.Kernel, chat_func: sk.SKFunctionBase, context: sk.SKContext) -> bool:\n",
    "    try:\n",
    "        user_input = input(\"用户:> \")\n",
    "        context[\"user_input\"] = user_input\n",
    "        print(f\"User:> {user_input}\")\n",
    "    except KeyboardInterrupt:\n",
    "        print(\"\\n\\nExiting chat...\")\n",
    "        return False\n",
    "    except EOFError:\n",
    "        print(\"\\n\\nExiting chat...\")\n",
    "        return False\n",
    "\n",
    "    if user_input == \"exit\":\n",
    "        print(\"\\n\\nExiting chat...\")\n",
    "        return False\n",
    "\n",
    "    print(context.variables)\n",
    "    \n",
    "    answer = await kernel.run_async(chat_func, input_vars=context.variables)\n",
    "    context[\"chat_history\"] += f\"\\n当前输入:> {user_input}\\n回答:> {answer}\\n\"\n",
    "\n",
    "    print(f\"Bot:> {answer}\")\n",
    "    return True"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "id": "e3875a34",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "注入背景...\n",
      "构建prompt...\n",
      "开始对话 (type 'exit' to exit):\n",
      "\n",
      "User:> exit\n",
      "\n",
      "\n",
      "Exiting chat...\n"
     ]
    }
   ],
   "source": [
    "#-# cell_skip\n",
    "print(\"注入背景...\")\n",
    "await populate_memory(kernel)\n",
    "\n",
    "# print(\"开始提问\")\n",
    "# await search_memory_examples(kernel)\n",
    "\n",
    "print(\"构建prompt...\")\n",
    "chat_func, context = await setup_chat_with_memory(kernel)\n",
    "\n",
    "print(\"开始对话 (type 'exit' to exit):\\n\")\n",
    "chatting = True\n",
    "while chatting:\n",
    "    chatting = await chat(kernel, chat_func, context)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "0a51542b",
   "metadata": {},
   "source": [
    "### 添加外部links到Memory中\n",
    "\n",
    "很多时候，我们有大量的外部知识库，接下来我们将使用SK的`VolatileMemoryStore`以用于加载外部链接：\n",
    "例如我们添加千帆SDK的repo："
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "id": "c3d5a1b9",
   "metadata": {},
   "outputs": [],
   "source": [
    "github_files = {\n",
    "    \"https://github.com/baidubce/bce-qianfan-sdk/blob/main/README.md\": \"README: 千帆SDK介绍，安装，基础使用方法\",\n",
    "    \"https://github.com/baidubce/bce-qianfan-sdk/blob/main/cookbook/finetune/trainer_finetune.ipynb\": \"Cookbook: 千帆SDK Trainer使用方法\"\n",
    "}"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "75f3ea5e",
   "metadata": {},
   "source": [
    "与之前不同，这里通过 `SaveReferenceAsync`将数据和引用来源分开存储"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "id": "170e7142",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "  已添加 1 saved\n",
      "  已添加 2 saved\n"
     ]
    }
   ],
   "source": [
    "memory_collection_name = \"random_QianfanGithub\"\n",
    "i = 0\n",
    "for entry, value in github_files.items():\n",
    "    await kernel.memory.save_reference_async(\n",
    "        collection=memory_collection_name,\n",
    "        description=value,\n",
    "        text=value,\n",
    "        external_id=entry,\n",
    "        external_source_name=\"GitHub\",\n",
    "    )\n",
    "    i += 1\n",
    "    print(\"  已添加 {} saved\".format(i))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 14,
   "id": "143911c3",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "===========================\n",
      "Query: 我希望整体了解千帆SDK，有什么办法？\n",
      "\n",
      "Result 1:\n",
      "  URL:     : https://github.com/baidubce/bce-qianfan-sdk/blob/main/README.md\n",
      "  Title    : README: 千帆SDK介绍，安装，基础使用方法\n",
      "  Relevance: 0.7502846678234273\n"
     ]
    }
   ],
   "source": [
    "ask = \"我希望整体了解千帆SDK，有什么办法？\"\n",
    "print(\"===========================\\n\" + \"Query: \" + ask + \"\\n\")\n",
    "\n",
    "results = await kernel.memory.search_async(memory_collection_name, ask, limit=5, min_relevance_score=0.7)\n",
    "\n",
    "i = 0\n",
    "for res in results:\n",
    "    i += 1\n",
    "    print(f\"Result {i}:\")\n",
    "    print(\"  URL:     : \" + res.id)\n",
    "    print(\"  Title    : \" + res.description)\n",
    "    print(\"  Relevance: \" + str(res.relevance))\n",
    "    print()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "3d43cdc1",
   "metadata": {},
   "source": [
    "使用返回的文档进行"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "59294dac",
   "metadata": {},
   "source": [
    "除了VolatileMemory之外，我们还可以通过对接外部向量库的形式实现大量的外部知识库，SK官方提供常用的例如`chroma`,`pinecone`等实现。通过直接替换memory_store可以实现kernel和chroma的对接："
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 15,
   "id": "82564625b959a3c0",
   "metadata": {},
   "outputs": [],
   "source": [
    "from semantic_kernel.connectors.memory.chroma import (\n",
    "    ChromaMemoryStore,\n",
    ")\n",
    "\n",
    "kernel.register_memory_store(\n",
    "    memory_store=ChromaMemoryStore(\n",
    "        persist_directory=\"./\"\n",
    "    )\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 16,
   "id": "7bcbca175c7b40d7",
   "metadata": {},
   "outputs": [],
   "source": [
    "await populate_memory(kernel)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 17,
   "id": "84563f63220180f0",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Question: 我的名字是？\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Chroma returns distance score not cosine similarity score.                So embeddings are automatically queried from database for calculation.\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Answer: 我名字叫做小度\n",
      "\n",
      "Question: 我在哪里工作？\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Chroma returns distance score not cosine similarity score.                So embeddings are automatically queried from database for calculation.\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Answer: 我工作在baidu\n",
      "\n",
      "Question: 我去过哪些地方旅游?\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Chroma returns distance score not cosine similarity score.                So embeddings are automatically queried from database for calculation.\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Answer: 我曾去过北京，上海，深圳\n",
      "\n",
      "Question: 我的家乡是?\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Chroma returns distance score not cosine similarity score.                So embeddings are automatically queried from database for calculation.\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Answer: 我来自中国\n",
      "\n",
      "Question: 我的爱好是？\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Chroma returns distance score not cosine similarity score.                So embeddings are automatically queried from database for calculation.\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Answer: 我爱打羽毛球\n"
     ]
    }
   ],
   "source": [
    "await search_memory_examples(kernel)"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3 (ipykernel)",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.9.13"
  },
  "vscode": {
   "interpreter": {
    "hash": "58f7cb64c3a06383b7f18d2a11305edccbad427293a2b4afa7abe8bfc810d4bb"
   }
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}
